package com.octopus.core.downloader;

import com.octopus.core.Request;
import com.octopus.core.Response;
import com.octopus.core.downloader.proxy.Proxy;
import com.octopus.core.utils.DownloaderUtils;
import com.octopus.core.utils.StringUtils;
import java.io.ByteArrayOutputStream;
import java.io.InputStream;
import java.net.CookieHandler;
import java.net.CookieManager;
import java.net.CookiePolicy;
import java.net.HttpCookie;
import java.net.HttpURLConnection;
import java.net.InetSocketAddress;
import java.net.Proxy.Type;
import java.net.URI;
import java.net.URISyntaxException;
import java.net.URL;
import java.security.KeyManagementException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.util.Base64;
import java.util.Map;
import java.util.zip.GZIPInputStream;
import javax.net.ssl.HostnameVerifier;
import javax.net.ssl.HttpsURLConnection;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLSocketFactory;
import javax.net.ssl.TrustManager;
import lombok.Cleanup;
import lombok.extern.slf4j.Slf4j;

/**
 * 默认页面下载器
 *
 * <p>支持基于用户名/密码基本认证的代理服务器
 *
 * @author shoulai.yang@gmail.com
 * @date 2019/07/31
 */
@Slf4j
public class DefaultDownloader implements Downloader {

  private HostnameVerifier hostnameVerifier = new TrustAllHostnameVerifier();

  private CookieManager manager = new CookieManager();

  private static final String GZIP = "gzip";

  {
    try {
      SSLContext sslContext = SSLContext.getInstance("TLS");
      sslContext.init(null, new TrustManager[] {new HttpsTrustManager()}, new SecureRandom());
      SSLSocketFactory sslSocketFactory = sslContext.getSocketFactory();
      HttpsURLConnection.setDefaultSSLSocketFactory(sslSocketFactory);
      HttpsURLConnection.setDefaultHostnameVerifier(hostnameVerifier);
      manager.setCookiePolicy(CookiePolicy.ACCEPT_ALL);
      CookieHandler.setDefault(manager);
    } catch (NoSuchAlgorithmException | KeyManagementException e) {
      e.printStackTrace();
    }
  }

  @Override
  public Response download(Request request, Proxy proxy) throws DownloadException {
    Response page = Response.me(request);
    HttpURLConnection connection = null;
    try {
      this.storeCookie(request);
      URL url = new URL(request.url());
      connection =
          (HttpURLConnection)
              url.openConnection(
                  proxy != null
                      ? new java.net.Proxy(
                          Type.HTTP, new InetSocketAddress(proxy.getHost(), proxy.getPort()))
                      : java.net.Proxy.NO_PROXY);
      connection.setRequestMethod(request.method());
      connection.setConnectTimeout(request.connectTimeout());
      connection.setReadTimeout(request.readTimeout());
      connection.setDoInput(true);
      connection.setDoOutput(true);
      // set header
      this.setHeader(request.headers(), connection);
      if (proxy != null && !StringUtils.isEmpty(proxy.getUsername())) {
        connection.addRequestProperty(
            "Proxy-Authorization",
            "Basic "
                + Base64.getEncoder()
                    .encodeToString((proxy.getUsername() + ":" + proxy.getPassword()).getBytes()));
      }
      connection.connect();
      int httpStatus = connection.getResponseCode();
      page.httpStatus(httpStatus);
      String contentEncoding = connection.getContentEncoding();
      @Cleanup InputStream inputStream = connection.getInputStream();
      if (contentEncoding != null && contentEncoding.toLowerCase().contains(GZIP)) {
        inputStream = new GZIPInputStream(inputStream);
      }
      page.charset(DownloaderUtils.getCharsetFromContentType(connection.getContentType()));
      @Cleanup ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
      // 10kb
      int bufferSize = 10240;
      byte[] buffer = new byte[bufferSize];
      int length;
      while ((length = inputStream.read(buffer, 0, bufferSize)) != -1) {
        outputStream.write(buffer, 0, length);
      }
      page.body(outputStream.toByteArray());
    } catch (Exception e) {
      throw new DownloadException(request, e);
    } finally {
      if (connection != null) {
        connection.disconnect();
      }
    }
    return page;
  }

  private void storeCookie(Request request) {
    try {
      Map<String, String> cookies = request.cookies();
      if (cookies != null && !cookies.isEmpty()) {
        URI uri = new URI(request.url());
        String host = uri.getHost();
        cookies.forEach(
            (key, value) -> {
              HttpCookie httpCookie = new HttpCookie(key, value);
              httpCookie.setDomain(host);
              httpCookie.setPath("/");
              httpCookie.setVersion(0);
              manager.getCookieStore().add(uri, httpCookie);
            });
      }
    } catch (URISyntaxException e) {
      e.printStackTrace();
    }
  }

  private void setHeader(Map<String, String> headers, HttpURLConnection connection) {
    if (headers != null && !headers.isEmpty()) {
      headers.forEach(connection::addRequestProperty);
    }
  }
}
